/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: 对拿到partition map的key进行去重，并将拿到的value恢复.
 * Author: qingjicheng
 * Create: 2022-11-28
 */


#ifndef DROP_DUPLICATES
#define DROP_DUPLICATES
#include <unordered_set>
#include <hccl_types.h>
#include <mutex>
#include "base.h"
#include "mem_device_pub.h"

namespace hccl {
template<typename KeyType>
class DropDuplicates {
public:
    struct RecoverParams {
        void *value = nullptr;
        u64 valueMemSize = 0;
        u32 valueItemSize = 0;
        KeyType *keys = nullptr;
        u32 keyMaxNum = 0;
        KeyType *rawKeys = nullptr;
        u32 rawKeyMaxNum = 0;
        void *tempMem = nullptr;
        u64 tempMemSize = 0;
    };

public:
    DropDuplicates() = default;
    ~DropDuplicates() = default;

    std::vector<KeyType> Drop(KeyType *keys, u32 keyMaxNum)
    {
        std::unordered_set<KeyType> keysSet(keys, keys + keyMaxNum);
        std::vector<KeyType> keysVec(keysSet.begin(), keysSet.end());
        return std::move(keysVec);
    }

    HcclResult Recover(RecoverParams params)
    {
        CHK_SAFETY_FUNC_RET(memcpy_s(params.tempMem, params.tempMemSize, params.value, params.valueMemSize));
        for (u32 i = 0; i < params.keyMaxNum; i++) {
            u8 *value = static_cast<u8 *>(params.tempMem) + i * params.valueItemSize;
            for (u32 k = 0; k < params.rawKeyMaxNum; k++) {
                if (*(params.keys + i) == *(params.rawKeys + k)) {
                    CHK_SAFETY_FUNC_RET(memcpy_s(static_cast<u8 *>(params.value) + k * params.valueItemSize,
                        params.valueMemSize - k * params.valueItemSize, value, params.valueItemSize));
                }
            }
        }
        return HCCL_SUCCESS;
    }
};
}
#endif // DROP_DUPLICATES